#include <stdio.h>
#include <stdlib.h>
#include <stdint.h>
#include <sys/types.h>
#include <sys/stat.h>
#include <signal.h>
#include <pthread.h>
#include <getopt.h>
#include <string.h>
#include <limits.h>


#include <rte_eal.h>
#include <rte_mbuf.h>
#include <rte_malloc.h>
#include <rte_log.h>
#include <rte_ethdev.h>
#include <rte_hash.h>
#include <rte_lcore.h>
#include <rte_string_fns.h>
#include <rte_vhost.h>

#include "switch.h"

#ifndef PATH_MAX
#define PATH_MAX 4096 
#endif


#define NUM_MBUFS 8192
#define MBUF_CACHE_SIZE 250  // 缓存大小 
#define MAX_PKT_BURST 32

#define RX_RING_SIZE 1024  // 接收环队列大小 
#define TX_RING_SIZE 1024  // 发送环队列大小 

#define NUM_HASH_ENTRIES 256

static struct rte_mempool *mbuf_pool;  // 数据包内存池
static int nb_ports = 0;               // 网卡端口数量
static int nb_sockets = 0;             // vhost套接字数量
static char *socket_files = NULL;      // 存储套接字文件路径

static struct dev_info_tailq dev_list; // 设备列表，用于存储设备信息
static pthread_mutex_t dev_list_mutex; // 设备列表的互斥锁
static volatile int dev_list_update;   // 标记设备列表是否更新
static volatile int force_exit = 0;    // 标记是否强制退出

static struct rte_hash *mac_output_map;      // MAC地址到输出端口的映射表
static struct dev_info *output_table[NUM_HASH_ENTRIES]; // 输出表，用于存储输出端口信息

/*
 * 显示命令行参数的用法
 */
static void
print_usage()
{
    printf("\nTo add each socket file, use parameter -s or --socket-file\n"
           "For example:\n"
           "./learning-switch -c 3 -n 4 -m 256 -- "
           "-s /PATH/TO/SOCKET/FILE1 "
           "--socket-file /PATH/TO/SOCKET/FILE2\n\n");
}

/*
 * 解析命令行参数，获取vhost套接字文件路径
 */
static int
parse_socket_paths(int argc, char **argv) {
    int opt;
    static struct option long_option[] = {
        { "help", no_argument, NULL, 'h'},
        { "socket-file", required_argument, NULL, 's'},
        { NULL, 0, 0, 0},
    };

    // 使用 getopt_long 来解析命令行参数
    while((opt = getopt_long(argc, argv, "hs:", long_option, NULL)) != -1) {
        switch (opt) {
        case 'h':  // 帮助信息
            print_usage();
            exit(0);
            break;
        case 's':  // 获取socket文件路径
            if (strlen(optarg) > PATH_MAX) {
                rte_log(RTE_LOG_INFO, RTE_LOGTYPE_USER1,
                        "Invalid path length for socket name (Max %d characters)\n", PATH_MAX);
                return -1;
            }
            // 动态分配套接字文件路径的存储空间
            socket_files = realloc(socket_files, PATH_MAX * (nb_sockets + 1));
            if (!socket_files) {
                rte_log(RTE_LOG_ERR, RTE_LOGTYPE_USER1, "Memory allocation error for socket files\n");
                return -1;
            }
            // 将路径存储到 socket_files
            snprintf(socket_files + nb_sockets * PATH_MAX, PATH_MAX, "%s", optarg);
            nb_sockets++;
            break;
        default:  // 其他无效参数
            rte_log(RTE_LOG_INFO, RTE_LOGTYPE_USER1, "Invalid arguments\n");
            print_usage();
            return -1;
        }
    }
    return 0;
}


static int
port_init(int portid)
{
    // 端口配置：仅设置最大接收包长度
    const struct rte_eth_conf port_conf = {
        .rxmode = { .max_rx_pkt_len = RTE_ETHER_MAX_LEN }
    };
    const uint16_t rx_rings = 1;  // 接收环队列数
    const uint16_t tx_rings = 1;  // 发送环队列数
    const uint16_t qid = 0;       // 队列ID
    struct dev_info *dev;
    int ret;

    /* 检查内存池是否创建成功 */
    if (mbuf_pool == NULL) {
        rte_log(RTE_LOG_ERR, RTE_LOGTYPE_USER1, "Error: Mbuf pool is not initialized\n");
        return -1;
    }

    /* 配置以太网设备 */
    ret = rte_eth_dev_configure(portid, rx_rings, tx_rings, &port_conf);
    if (ret != 0) {
        rte_log(RTE_LOG_ERR, RTE_LOGTYPE_USER1, "Error configuring Ethernet device: port=%d, ret=%d\n", portid, ret);
        return ret;
    }

    /* 设置接收队列 */
    ret = rte_eth_rx_queue_setup(portid, qid, RX_RING_SIZE, rte_eth_dev_socket_id(portid), NULL, mbuf_pool);
    if (ret < 0) {
        rte_log(RTE_LOG_ERR, RTE_LOGTYPE_USER1, "Error setting up RX queue: port=%d, ret=%d\n", portid, ret);
        return ret;
    }

    /* 设置发送队列 */
    ret = rte_eth_tx_queue_setup(portid, qid, TX_RING_SIZE, rte_eth_dev_socket_id(portid), NULL);
    if (ret < 0) {
        rte_log(RTE_LOG_ERR, RTE_LOGTYPE_USER1, "Error setting up TX queue: port=%d, ret=%d\n", portid, ret);
        return ret;
    }

    /* 启动以太网端口 */
    ret = rte_eth_dev_start(portid);
    if (ret < 0) {
        rte_log(RTE_LOG_ERR, RTE_LOGTYPE_USER1, "Error starting Ethernet port: port=%d, ret=%d\n", portid, ret);
        return ret;
    }

    /* 获取并显示端口的MAC地址 */
    struct rte_ether_addr addr;
    rte_eth_macaddr_get(portid, &addr);
    rte_log(RTE_LOG_INFO, RTE_LOGTYPE_USER1, "Port %u MAC: "
            "%02" PRIx8 ":%02" PRIx8 ":%02" PRIx8 ":"
            "%02" PRIx8 ":%02" PRIx8 ":%02" PRIx8 "\n",
            (unsigned)portid,
            addr.addr_bytes[0], addr.addr_bytes[1],
            addr.addr_bytes[2], addr.addr_bytes[3],
            addr.addr_bytes[4], addr.addr_bytes[5]);

    /* 启用接收端口的混杂模式 */
    rte_eth_promiscuous_enable(portid);

    /* 分配并初始化设备信息结构体 */
    dev = (struct dev_info *)rte_zmalloc("device info", sizeof(*dev), RTE_CACHE_LINE_SIZE);
    if (dev == NULL) {
        rte_log(RTE_LOG_ERR, RTE_LOGTYPE_USER1, "Error allocating memory for device info: port=%d\n", portid);
        return -1;
    }
    dev->virtual = 0;
    dev->id = portid;
    dev->state = DEVICE_READY;

    /* 将设备信息加入设备列表 */
    pthread_mutex_lock(&dev_list_mutex);
    TAILQ_INSERT_TAIL(&dev_list, dev, dev_entry);
    pthread_mutex_unlock(&dev_list_mutex);

    rte_log(RTE_LOG_INFO, RTE_LOGTYPE_USER1, "Physical port %d added to the device list\n", dev->id);

    return 0;
}


/*
 * 当检测到新的虚拟设备时调用的回调函数
 */
static int
new_vdev_callback(int vid)
{
    struct dev_info *dev;

    /* 分配并初始化虚拟设备信息结构体 */
    dev = (struct dev_info *)rte_zmalloc("vhost device", sizeof(*dev), RTE_CACHE_LINE_SIZE);
    if (dev == NULL) {
        rte_log(RTE_LOG_ERR, RTE_LOGTYPE_USER1, "Error allocating memory for virtual device: vid=%d\n", vid);
        return -1;
    }

    /* 初始化设备信息 */
    dev->virtual = 1;         // 标记为虚拟设备
    dev->id = vid;            // 设置虚拟设备ID
    dev->state = DEVICE_READY; // 设备状态设置为“已就绪”

    /* 通知转发线程解锁设备列表 */
    dev_list_update = 1;

    /* 加锁以确保线程安全 */
    pthread_mutex_lock(&dev_list_mutex);

    /* 将新虚拟设备添加到设备列表末尾 */
    TAILQ_INSERT_TAIL(&dev_list, dev, dev_entry);

    /* 重置更新标记并解锁 */
    dev_list_update = 0;
    pthread_mutex_unlock(&dev_list_mutex);

    /* 日志记录虚拟设备添加操作 */
    rte_log(RTE_LOG_INFO, RTE_LOGTYPE_USER1,
            "Virtual device %d added to the device list\n", dev->id);

    return 0;
}


/*
 * 当虚拟设备被移除时调用的回调函数
 */
static void
destroy_vdev_callback(int vid)
{
    struct dev_info *dev = NULL;
    struct rte_ether_addr *key;
    int ret;

    /* 在设备列表中找到要删除的虚拟设备 */
    TAILQ_FOREACH(dev, &dev_list, dev_entry) {
        if (!dev->virtual)  // 跳过非虚拟设备
            continue;
        if (dev->id == vid)
            break;  // 找到匹配的设备
    }

    /* 如果设备未找到，则直接返回 */
    if (dev == NULL) {
        rte_log(RTE_LOG_INFO, RTE_LOGTYPE_USER1, "Device with ID %d not found in the list.\n", vid);
        return;
    }

    /* 设置设备状态为 DEVICE_CLOSING，通知转发线程停止交互 */
    dev->state = DEVICE_CLOSING;

    /* 通知转发线程解锁设备列表 */
    dev_list_update = 1;

    /* 加锁以确保线程安全 */
    pthread_mutex_lock(&dev_list_mutex);

    /* 从设备列表中移除该设备 */
    TAILQ_REMOVE(&dev_list, dev, dev_entry);

    /* 从哈希表中删除所有与该设备关联的 MAC 映射 */
    for (uint32_t i = 0; i < NUM_HASH_ENTRIES; i++) {
        if (output_table[i] != NULL && output_table[i]->virtual &&
            output_table[i]->id == vid) {
            
            /* 获取哈希表中位置 i 的键 */
            ret = rte_hash_get_key_with_position(mac_output_map, i, (void **)&key);
            if (ret < 0 || key == NULL) {
                rte_log(RTE_LOG_ERR, RTE_LOGTYPE_USER1, 
                        "Error retrieving MAC key for virtual device ID %d at position %d\n", vid, i);
                continue;
            }

            /* 从哈希表中删除该键 */
            ret = rte_hash_del_key(mac_output_map, key);
            if (ret < 0) {
                rte_log(RTE_LOG_ERR, RTE_LOGTYPE_USER1,
                        "Couldn't delete MAC: %02" PRIx8 ":%02" PRIx8 ":%02" PRIx8 ":%02" PRIx8 ":%02" PRIx8 ":%02" PRIx8
                        " from the hash table\n",
                        key->addr_bytes[0], key->addr_bytes[1], key->addr_bytes[2],
                        key->addr_bytes[3], key->addr_bytes[4], key->addr_bytes[5]);
            } else {
                rte_log(RTE_LOG_INFO, RTE_LOGTYPE_USER1, "MAC: %02" PRIx8 ":%02" PRIx8 ":%02" PRIx8 ":%02" PRIx8 ":%02" PRIx8 ":%02" PRIx8
                        " is removed from the hash table\n",
                        key->addr_bytes[0], key->addr_bytes[1], key->addr_bytes[2],
                        key->addr_bytes[3], key->addr_bytes[4], key->addr_bytes[5]);
                output_table[i] = NULL;  // 清除 output_table 中的对应条目
            }
        }
    }

    /* 解锁设备列表更新标志 */
    dev_list_update = 0;
    pthread_mutex_unlock(&dev_list_mutex);

    /* 记录虚拟设备已成功从设备列表中删除 */
    rte_log(RTE_LOG_INFO, RTE_LOGTYPE_USER1, "Virtual device %d deleted from the device list\n", dev->id);

    /* 释放设备结构体内存 */
    rte_free(dev);
}


/*
 * 哈希函数：将 MAC 地址映射到一个 8 位哈希值
 * 使用移位和异或操作来增加哈希值的分布均匀性
 */
static uint32_t
mac_to_uchar(const void *data, uint32_t data_len, uint32_t init_val)
{
    const uint8_t *d = (const uint8_t *)data;

    // 使用异或和移位操作来产生分布均匀的哈希值
    for (uint32_t i = 0; i < data_len; i++) {
        init_val ^= d[i];
        init_val = (init_val << 1) | (init_val >> 7);  // 左移一位并循环移位
    }

    // 只返回低8位值，确保返回结果在 8 位范围内
    return init_val & 0xFF;
}


/*
 * 从接收到的数据包中学习源 MAC 地址。该地址作为键存储在 mac-to-output 哈希表中，
 * 与之关联的值是 output_table 的索引，每个元素指向一个输出端口。
 */
static int
learn_mac_address(struct rte_mbuf *mbuf, struct dev_info *s_dev)
{
    struct rte_ether_hdr *pkt_hdr;
    int ret;

    /* 获取数据包的以太网头部并提取源 MAC 地址 */
    pkt_hdr = rte_pktmbuf_mtod(mbuf, struct rte_ether_hdr *);

    /* 将源 MAC 地址添加到 mac-to-output 哈希表中 */
    ret = rte_hash_add_key(mac_output_map, &pkt_hdr->s_addr);
    if (ret < 0) {
        rte_log(RTE_LOG_INFO, RTE_LOGTYPE_USER1,
                "Cannot add entry to mac-output lookup table (error %d)\n", ret);
        return -1;
    }

    /* 检查并更新 output_table 中与 MAC 地址关联的设备信息 */
    if (output_table[ret] == NULL ||
        output_table[ret]->virtual != s_dev->virtual ||
        output_table[ret]->id != s_dev->id) {

        /* 更新 output_table 的对应条目，指向当前设备信息 */
        output_table[ret] = s_dev;

        /* 记录成功学习 MAC 地址的日志信息 */
        rte_log(RTE_LOG_INFO, RTE_LOGTYPE_USER1, "Learned MAC: "
                "%02" PRIx8 ":%02" PRIx8 ":%02" PRIx8 ":%02" PRIx8 ":%02" PRIx8 ":%02" PRIx8
                " mapped to %s device with id %d\n",
                pkt_hdr->s_addr.addr_bytes[0], pkt_hdr->s_addr.addr_bytes[1],
                pkt_hdr->s_addr.addr_bytes[2], pkt_hdr->s_addr.addr_bytes[3],
                pkt_hdr->s_addr.addr_bytes[4], pkt_hdr->s_addr.addr_bytes[5],
                (s_dev->virtual) ? "virtual" : "physical", s_dev->id);
    }
    return 0;
}



/*
 * 根据以下条件转发给定的数据包：
 * 1) 如果数据包的目标地址不可识别，则将数据包广播到除传入设备外的所有端口/设备。
 * 2) 否则，根据目标地址执行单播转发。
 */
static void
forward_packet(struct rte_mbuf *mbuf, struct dev_info *s_dev)
{
    struct rte_ether_hdr *pkt_hdr;
    struct dev_info *dev;
    struct rte_mbuf *tbuf;
    int ret;

    /* 获取以太网头部信息，并查找目的地址的输出设备 */
    pkt_hdr = rte_pktmbuf_mtod(mbuf, struct rte_ether_hdr *);
    ret = rte_hash_lookup(mac_output_map, &pkt_hdr->d_addr);

    /* 广播处理：当目的地址不可识别时 */
    if (ret < 0) {
        TAILQ_FOREACH(dev, &dev_list, dev_entry) {
            /* 跳过数据包的来源设备 */
            if (dev == s_dev)
                continue;

            if (dev->virtual) {
                /* 检查虚拟设备是否处于关闭状态 */
                if (unlikely(dev->state == DEVICE_CLOSING))
                    continue;

                /* 将数据包发送到虚拟设备的接收队列 */
                rte_vhost_enqueue_burst(dev->id, 0, &mbuf, 1);
            } else {
                /* 克隆数据包以便发送给物理设备 */
                tbuf = rte_pktmbuf_clone(mbuf, mbuf_pool);
                if (tbuf == NULL) {
                    rte_log(RTE_LOG_ERR, RTE_LOGTYPE_USER1,
                            "Failed to clone packet for broadcast\n");
                    continue;
                }

                /* 将数据包发送到物理端口，若发送失败则释放克隆的包 */
                ret = rte_eth_tx_burst(dev->id, 0, &tbuf, 1);
                if (unlikely(ret == 0))
                    rte_pktmbuf_free(tbuf);
            }
        }

        /* 广播完成后释放原始数据包 */
        rte_pktmbuf_free(mbuf);
        return;
    }

    /* 单播处理：找到目标设备并转发 */
    dev = output_table[ret];
    if (dev->virtual) {
        /* 如果目标是虚拟设备且未处于关闭状态，转发数据包 */
        if (unlikely(dev->state != DEVICE_CLOSING))
            rte_vhost_enqueue_burst(dev->id, 0, &mbuf, 1);
        rte_pktmbuf_free(mbuf);
    } else {
        /* 如果目标是物理设备，发送数据包到物理端口 */
        ret = rte_eth_tx_burst(dev->id, 0, &mbuf, 1);
        if (unlikely(ret == 0)) {
            rte_log(RTE_LOG_ERR, RTE_LOGTYPE_USER1,
                    "Failed to send packet to physical device ID %d\n", dev->id);
            rte_pktmbuf_free(mbuf);
        }
    }
}



/*
 * 注销套接字文件并停止物理设备的函数
 */
static void cleanup(void)
{
    /* 注销每个 vhost 套接字文件 */
    for (int vid = 0; vid < nb_sockets; vid++) {
        int ret = rte_vhost_driver_unregister(&socket_files[vid * PATH_MAX]);
        if (ret < 0) {
            rte_log(RTE_LOG_ERR, RTE_LOGTYPE_USER1,
                    "Failed to unregister vhost socket file: %s\n",
                    &socket_files[vid * PATH_MAX]);
        } else {
            rte_log(RTE_LOG_INFO, RTE_LOGTYPE_USER1,
                    "Successfully unregistered vhost socket file: %s\n",
                    &socket_files[vid * PATH_MAX]);
        }
    }

    /* 停止每个物理设备 */
    for (int portid = 0; portid < nb_ports; portid++) {
        rte_eth_dev_stop(portid);
        rte_eth_dev_close(portid);
        rte_log(RTE_LOG_INFO, RTE_LOGTYPE_USER1,
                "Stopped and closed physical device with port id: %d\n", portid);
    }

    /* 销毁设备列表的互斥锁 */
    pthread_mutex_destroy(&dev_list_mutex);
    rte_log(RTE_LOG_INFO, RTE_LOGTYPE_USER1, "Destroyed device list mutex\n");

    /* 打印日志并退出程序 */
    rte_log(RTE_LOG_INFO, RTE_LOGTYPE_USER1, "Cleanup completed. Exiting...\n");
    exit(0);
}

/*
 * 主转发线程：大部分时间用于转发数据包，只有在需要更新设备列表时才会暂停转发。
 */
static int learning_switch_main(void *arg __rte_unused)
{
    struct rte_mbuf *mbufs[MAX_PKT_BURST];
    uint16_t nb_rcv;
    struct dev_info *dev;

    /* 主循环：检查 Force_exit 标志用于退出 */
    while (!force_exit) {
        /* 加锁以确保设备列表的安全访问 */
        pthread_mutex_lock(&dev_list_mutex);

        /* 在不需要更新设备列表且未收到退出信号时，执行转发操作 */
        while (!dev_list_update && !force_exit) {
            /* 遍历设备列表并处理每个设备的数据包 */
            TAILQ_FOREACH(dev, &dev_list, dev_entry) {
                if (dev->state == DEVICE_CLOSING)  // 跳过正在关闭的设备
                    continue;

                /* 根据设备类型接收数据包 */
                if (dev->virtual) {
                    nb_rcv = rte_vhost_dequeue_burst(dev->id, 0, mbuf_pool, mbufs, MAX_PKT_BURST);
                } else {
                    nb_rcv = rte_eth_rx_burst(dev->id, 0, mbufs, MAX_PKT_BURST);
                }

                /* 处理接收的数据包：学习源 MAC 地址并转发 */
                for (int n = 0; n < nb_rcv; n++) {
                    learn_mac_address(mbufs[n], dev);   // 学习 MAC 地址
                    forward_packet(mbufs[n], dev);      // 转发数据包
                }
            }
        }

        /* 释放锁，以允许其他线程更新设备列表 */
        pthread_mutex_unlock(&dev_list_mutex);
    }

    /* 清理资源并退出 */
    cleanup();
    return 0;
}


/*
 * 处理 SIGINT 信号的回调函数，用于优雅终止程序。
 */
static void force_exit_handling(__rte_unused int signum)
{
    force_exit = 1;
}


int main(int argc, char **argv)
{
    int ret;
    mode_t sock_perm;

    /* 初始化 DPDK 环境抽象层 (EAL) */
    ret = rte_eal_init(argc, argv);
    if (ret < 0)
        rte_exit(EXIT_FAILURE, "Error with EAL initialization\n");

    argc -= ret;
    argv += ret;

    /* 解析自定义的 socket 文件路径参数 */
    ret = parse_socket_paths(argc, argv);
    if (ret < 0)
        rte_exit(EXIT_FAILURE, "Invalid switch arguments\n");

    /* 注册 SIGINT 信号处理函数，用于终止和清理 */
    signal(SIGINT, force_exit_handling);

    /* 创建 mbuf 内存池，用于存储数据包 */
    nb_ports = rte_eth_dev_count_avail();  // 获取可用的物理端口数量
    mbuf_pool = rte_pktmbuf_pool_create("MBUF_POOL",
                                        NUM_MBUFS * (nb_ports + nb_sockets),
                                        MBUF_CACHE_SIZE, 0,
                                        RTE_MBUF_DEFAULT_BUF_SIZE,
                                        rte_socket_id());
    if (mbuf_pool == NULL)
        rte_exit(EXIT_FAILURE, "Cannot create mbuf pool\n");

    /* 创建 MAC 地址到输出端口的哈希表 */
    const struct rte_hash_parameters mac_output_map_params = {
        .name = "mac_output_map",
        .entries = NUM_HASH_ENTRIES,
        .key_len = sizeof(struct rte_ether_addr),
        .hash_func = mac_to_uchar,
        .hash_func_init_val = 0,
    };
    mac_output_map = rte_hash_create(&mac_output_map_params);
    if (mac_output_map == NULL)
        rte_exit(EXIT_FAILURE, "Cannot create MAC-output hash table\n");
    
    /* 初始化输出端口表 */
    memset(output_table, 0, sizeof(struct dev_info *) * NUM_HASH_ENTRIES);

    /* 初始化设备列表 */
    TAILQ_INIT(&dev_list);
    dev_list_update = 0;
    if (pthread_mutex_init(&dev_list_mutex, NULL) != 0) {
        rte_exit(EXIT_FAILURE, "Mutex initialization failure\n");
    }

    /* 初始化物理端口 */
    for (int portid = 0; portid < nb_ports; portid++) {
        if (port_init(portid) != 0) {
            rte_exit(EXIT_FAILURE,
                     "Physical port %d initialization failure\n",
                     portid);
        }
    }

    /* 设置 socket 文件权限为 0666 */
    sock_perm = umask(~0666);

    /* 注册 vhost 驱动 */
    for (int i = 0; i < nb_sockets; i++) {
        const char *socket_path = &socket_files[i * PATH_MAX];
        
        ret = rte_vhost_driver_register(socket_path, 0);
        if (ret != 0) {
            rte_exit(EXIT_FAILURE, "Vhost register failure for socket: %s\n",
                     socket_path);
        }

        /* 注册回调函数用于添加和删除虚拟设备 */
        const struct vhost_device_ops virtio_net_device_ops = {
            .new_device = new_vdev_callback,
            .destroy_device = destroy_vdev_callback,
        };

        ret = rte_vhost_driver_callback_register(socket_path, &virtio_net_device_ops);
        if (ret != 0) {
            rte_exit(EXIT_FAILURE, "Callback register failure for socket: %s\n",
                     socket_path);
        }
    }

    /* 将进程权限恢复到原始值 */
    umask(sock_perm);

    /* 启动转发线程 */
    rte_eal_remote_launch(learning_switch_main, NULL, 1);

    /* 启动 vhost 会话，用于处理虚拟设备的连接 */
    // rte_vhost_driver_session_start();

    /* 主线程等待转发线程完成 */
    rte_eal_mp_wait_lcore();
    
    /* 释放资源并优雅退出 */
    cleanup();
    return 0;
}
